# names for regression tables and plot labels
cmap = c(# main vars for main models
  "(Intercept)" = "Intercept",
  "treatTRUE" = "Treated",
  "running" = "Days since event",
  "treatTRUE:running" = "Treated $\\times$ days since event"
)

# get functions for modelsummary of rdrobust
tidy.rdrobust <- function(object, ...){
  ret <- data.frame(term = row.names(object$coef),
                    estimate = object$coef[, 1],
                    std.error = object$se[, 1],
                    statistic = object$z[, 1],
                    p.value = object$pv[, 1],
                    conf.low = object$ci[,1],
                    conf.high = object$ci[, 2])
  row.names(ret) <- NULL
  ret
}

glance.rdrobust <- function(object, ...){
  ret <- data.frame(nobs.left = object$N[1],
                    nobs.right = object$N[2],
                    nobs.effective.left = object$N_h[1],
                    nobs.effective.right = object$N_h[2],
                    cutoff = object$c,
                    order.regression = object$p,
                    order.bias = object$q,
                    kernel = object$kernel,
                    bwselect = object$bwselect,
                    bw.regression = round(object$bws[1], 2),
                    bw.bias = round(object$bws[2], 2))
  names(ret) = c("N (pre)", "N (post)", "Eff. N (pre)", "Eff. N (post)", "Cutoff",
                 "Order (reg.)", "Order (bias)", "Kernel", "Bw selection", "Bw (reg.)", "Bw (bias)")
  ret
}

# generate main RD plot
RDDplot = function(data, dv, running, treat, facet_var = NULL) {
  # fit model
  if (!is.null(facet_var)) {
    formula = as.formula(paste(dv, "~", running, "*", treat, "*", facet_var))
    data = data %>% filter(!is.na(!!sym(facet_var)))
  } else {
    formula = as.formula(paste(dv, "~", running, "*", treat))
  }
  fit <- lm(formula, data = data)

  # prediction grid
  if (!is.null(facet_var)) {
    pdat = expand.grid(running = seq(-6.5, 8.5, length.out = 200),
                       na.omit(unique(data[[facet_var]])))
    names(pdat)[2] = facet_var
  } else {
    pdat = expand.grid(running = seq(-6.5, 8.5, length.out = 200))
  }

  # cut of running around threshold
  pdat[[treat]] = case_when(
    pdat$running <= -0.5 ~ FALSE,
    pdat$running >= 0.5 ~ TRUE,
    TRUE ~ NA
  )

  pred = predict(fit, newdata = pdat, se.fit = TRUE)
  pdat$est = pred$fit
  pdat$upper = pdat$est + 1.96 * pred$se.fit
  pdat$lower = pdat$est - 1.96 * pred$se.fit

  # aggregate points
  pdat2 = data %>%
    group_by(!!sym(running), !!(if (!is.null(facet_var)) sym(facet_var) else NULL)) %>%
    summarize(n = log(sum(!is.na(!!sym(dv)))),
              avg_dv = mean(!!sym(dv), na.rm = TRUE), .groups = "drop")

  # Generate plot
  p = ggplot(pdat2, aes(x = !!sym(running), y = avg_dv, size = n)) +
    geom_vline(xintercept = 0, color = "grey70") +
    geom_ribbon(data = pdat,
                inherit.aes = FALSE,
                aes(x = running, ymin = lower, ymax = upper),
                fill = "gray90", alpha = 0.5, color = "black", linetype = 2) +
    geom_point(show.legend = FALSE,
               color = "gray30") +
    scale_size_continuous(range = c(.1, 2)) +
    geom_line(data = pdat,
              inherit.aes = FALSE,
              aes(x = running, y = est), color = "black") +
    scale_x_continuous(limits = c(-7, 9), breaks = seq(-7,9,2)) +
    theme_bw() +
    theme(strip.background = element_blank(),
          panel.grid = element_blank(),
          strip.text = element_text(face = "bold", size = 11)) +
    labs(y = "Propensity to vote for Greens (centered)",
         x = "Days since Council vote")

  # Add facets if facet_var is provided
  if (!is.null(facet_var)) {
    p = p + facet_wrap(as.formula(paste("~", facet_var)))
  }

  return(p)
}

# get legend from ggplot
g_legend <- function(a.gplot){
  tmp <- ggplot_gtable(ggplot_build(a.gplot))
  leg <- which(sapply(tmp$grobs, function(x) x$name) == "guide-box")
  legend <- tmp$grobs[[leg]]
  return(legend)}
